# -*- coding: utf-8 -*-
# @Time : 2021/6/7 下午5:43
# @Author : fugang_le
import time
import traceback
import sys
sys.path.append("/data/lefugang/template1/predict/template1")
from src.es_utils import ES_Util
from src.utils.util import time_difference
from src.utils.log_util import logger
from src.similar_engine.rerank import rerank, rerank_batch
from src.preprocess import input_preprocess
from flask_web.error_code import User
from configs.config import Facility_INDEX, MATERIAL_INDEX
from configs.classification_config import FACILITY_RESOURCE_POOL, MATERIAL_RESOURCE_POOL

GENERAL_MATCT = "MATCHT"    # 常规匹配
KEY_MATCH = "KEY_MATCH"     # 关键词检索
ALL_MATCH = 'ALL_MATCH'     # 全库匹配


def get_es_index(price_type):
    if price_type == FACILITY_RESOURCE_POOL:
        return Facility_INDEX
    if price_type == MATERIAL_RESOURCE_POOL:
        return MATERIAL_INDEX
    raise ValueError("es index error!")


def duplication(data, threshold):
    return [item['_source'] for item in data if item['_source']['score'] >= threshold]


def identification_process(price_type, grid_type, classification, materialName, sepc, price_conditions, threshold=0, topk=100):
    materialName = input_preprocess(materialName, classification)
    sepc = input_preprocess(sepc, classification)
    index = get_es_index(price_type)
    # set es size
    size = max(topk+100, 500)
    RETRIEVAL_TYPE = GENERAL_MATCT
    try:
        es_start_time = time.time()
        candidate_result = ES_Util.search_price_condition(index, materialName + sepc, classification, grid_type, price_conditions, size)
        # 判断是否要关键词检索
        if not candidate_result:
            # 不需要分类
            candidate_result = ES_Util.search_price_condition_by_wildcard(index, materialName + sepc, classification, grid_type, price_conditions, size)
            RETRIEVAL_TYPE = KEY_MATCH
        # if not candidate_result:
        #     candidate_result = ES_Util.search_price_condition(index, materialName + sepc, '', grid_type, price_conditions, size)
        #     RETRIEVAL_TYPE = ALL_MATCH
        logger.info("es search spend time:{}ms".format(time_difference(es_start_time)))
    except Exception:
        traceback.print_exc()
        return User.SEARCH
    print(candidate_result)
    rerank_start_time = time.time()
    query_dict = {'materialName': materialName, 'sepc': sepc}
    result = rerank_batch(query_dict, candidate_result)
    logger.info("rerank spend time:{}ms".format(time_difference(rerank_start_time)))
    # output处理, 去重处理
    if RETRIEVAL_TYPE == GENERAL_MATCT:
        result = duplication(result, threshold)
    # 关键词检索返回 top10
    elif RETRIEVAL_TYPE == KEY_MATCH or RETRIEVAL_TYPE == ALL_MATCH:
        result = duplication(result, 0)
        topk = 10
    return result[:topk]


def classify_identify_process(price_type, grid_type, classification, materialName, sepc, price_conditions, threshold=0, topk=100):
    materialName = input_preprocess(materialName, classification)
    sepc = input_preprocess(sepc, classification)
    index = get_es_index(price_type)
    # set es size
    size = max(topk+100, 500)
    RETRIEVAL_TYPE = GENERAL_MATCT
    try:
        es_start_time = time.time()
        candidate_result = ES_Util.search_price_condition(index, materialName + sepc, classification, grid_type, price_conditions, size)
        # 判断是否要关键词检索
        if not candidate_result:
            # 不需要分类
            candidate_result = ES_Util.search_price_condition_by_wildcard(index, materialName + sepc, '', grid_type, price_conditions, size)
            RETRIEVAL_TYPE = KEY_MATCH
        if not candidate_result:
            candidate_result = ES_Util.search_price_condition(index, materialName + sepc, '', grid_type, price_conditions, size)
            RETRIEVAL_TYPE = ALL_MATCH
        logger.info("es search spend time:{}ms".format(time_difference(es_start_time)))
    except Exception:
        traceback.print_exc()
        return User.SEARCH
    print(candidate_result)
    rerank_start_time = time.time()
    query_dict = {'materialName': materialName, 'sepc': sepc}
    result = rerank_batch(query_dict, candidate_result)
    logger.info("rerank spend time:{}ms".format(time_difference(rerank_start_time)))
    # output处理, 去重处理
    if RETRIEVAL_TYPE == GENERAL_MATCT:
        result = duplication(result, threshold)
    # 关键词检索返回 top10
    elif RETRIEVAL_TYPE == KEY_MATCH or RETRIEVAL_TYPE == ALL_MATCH:
        result = duplication(result, 0)
        topk = 10
    return result[:topk]


if __name__ == '__main__':
    price_conditions = {
        'price_condition_one': {'fir_area': '国网', 'seco_area': '总部,福建', 'third_area': '', 'particular_year': '2021',
                                'time_version': '二季度'},
        'price_condition_two': {'fir_area': '国网', 'seco_area': '总部,福建', 'third_area': '', 'particular_year': '2020',
                                'time_version': '一季度'},
        'price_condition_three': {'fir_area': '', 'seco_area': '', 'third_area': '', 'particular_year': '',
                                  'time_version': ''},}
    # price_conditions = {
    #     'price_condition_one': {'fir_area': '', 'seco_area': '', 'third_area': '', 'particular_year': '',
    #                             'time_version': ''},
    #     'price_condition_two': {'fir_area': '', 'seco_area': '', 'third_area': '', 'particular_year': '',
    #                             'time_version': ''},
    #     'price_condition_three': {'fir_area': '', 'seco_area': '', 'third_area': '', 'particular_year': '',
    #                               'time_version': ''},}
    res = process("设备资源库", '全部', "", "ztests中性点成套装置", '规格kg', price_conditions, 0.5, 100)
    for r in res:
        # print(r['materialName'], r['sepc'],r['score'])
        print(r)
        print()

